# -*- coding: utf-8 -*-

import pymysql


# 数据库操作函数，支持查询和更新
class MysqlClient(object):
    """
     with MysqlConnection('localhost', 'user', 'password', 'dbname') as db:
        db.insert("INSERT INTO your_table (name, age) VALUES (%s, %s)", ('Alice', 30))
        ......
    """

    def __init__(self, host: str, user: str, password: str, database: str, port=3306):
        self.host = host
        self.port = port
        self.user = user
        self.password = password
        self.database = database
        self.connection = None
        self.cursor = None

    def __enter__(self):  # 只有类中存在__enter__方法和__exit__方法，才可以调用with语句
        self.connection = pymysql.connect(
            host=self.host,
            port=self.port,
            user=self.user,
            password=self.password,
            database=self.database,
            charset='utf8mb4',
        )
        self.cursor = self.connection.cursor()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        if self.cursor:
            self.cursor.close()
        if self.connection:
            self.connection.close()

    def execute_sql(self, sql: str, params=()):
        try:
            if params:
                self.cursor.execute(sql, params)
            else:
                self.cursor.execute(sql)

            # 提交事务
            self.connection.commit()

            # 判断 SQL 语句类型
            if sql.strip().upper().startswith("SELECT"):
                return self.cursor.fetchall()  # 返回查询结果
            else:
                return self.cursor.rowcount, self.cursor.lastrowid  # 增、删、改，返回影响的行数和最后插入id
        except Exception as e:
            print("An error occurred:", e)
            self.connection.rollback()  # 出现错误时回滚事务
            return None
